#include "controller_base.h"

ControllerBase::ControllerBase(){
    speedPub_ = nh_.advertise<geometry_msgs::Twist>("/cmd_vel", 10);
    // path_planner_ptr_ = boost::make_shared<JPSPlanner>();
    path_planner_ptr_ = boost::make_shared<AstarPlanner>();
    origin_trajectory_.clear();
    baseFrameId_ = "base_link";
}

/*************************************************
* 数据同步
*************************************************/
// 同步代价地图 
void ControllerBase::updateCostmapData(boost::shared_ptr<nav_msgs::OccupancyGrid> costmap_ptr){
    map_base_.setMap(costmap_ptr);
    path_planner_ptr_->setMap(costmap_ptr);
}

// 同步机器人位姿
void ControllerBase::updateRobotPose(geometry_msgs::PoseStamped robot_pose){
    robot_pose_ = trajectoryROS_.toTrajectoryPoint2D(robot_pose);
}

// 同步原始全局路径
void ControllerBase::updateGlobalPath(vector<geometry_msgs::PoseStamped> path){
    origin_trajectory_.resize(path.size());
    for(int i=0; i<path.size(); i++){
        origin_trajectory_[i] = trajectoryROS_.toTrajectoryPoint2D(path[i]);
    }

    trajectory_points_ = origin_trajectory_;
}

// 获取实时全局路径 
vector<geometry_msgs::PoseStamped> ControllerBase::getGlobalPath(){
    vector<geometry_msgs::PoseStamped> path;
    for(int i=0; i<trajectory_points_.size(); i++){
        path.emplace_back(trajectoryROS_.toROSPoseStamped(trajectory_points_[i]));
    }
    return path;
}

/*************************************************
* 位姿转换
*************************************************/
// 获取TF位姿
bool ControllerBase::getTFPose(tf2_ros::Buffer* tf_buffer, geometry_msgs::PoseStamped& global_pose, 
string father_frame, string child_frame, ros::Time stamp){
    tf2::toMsg(tf2::Transform::getIdentity(), global_pose.pose);
    geometry_msgs::PoseStamped robot_pose;
    tf2::toMsg(tf2::Transform::getIdentity(), robot_pose.pose);
    robot_pose.header.frame_id = child_frame;
    robot_pose.header.stamp = ros::Time();
    // get the global pose of the robot
    try
    {
        // use current time if possible (makes sure it's not in the future)
        if (tf_buffer->canTransform(father_frame, child_frame, stamp))
        {
            geometry_msgs::TransformStamped transform = tf_buffer->lookupTransform(father_frame, child_frame, stamp);
            tf2::doTransform(robot_pose, global_pose, transform);
        }
        // use the latest otherwise
        else
        {
            tf_buffer->transform(robot_pose, global_pose, father_frame);
        }
    }
    catch (tf2::LookupException& ex)
    {
        ROS_ERROR_THROTTLE(1.0, "No Transform available Error looking up robot pose: %s\n", ex.what());
        return false;
    }
    catch (tf2::ConnectivityException& ex)
    {
        ROS_ERROR_THROTTLE(1.0, "Connectivity Error looking up robot pose: %s\n", ex.what());
        return false;
    }
    catch (tf2::ExtrapolationException& ex)
    {
        ROS_ERROR_THROTTLE(1.0, "Extrapolation Error looking up robot pose: %s\n", ex.what());
        return false;
    }
    // // check global_pose timeout
    // ros::Time current_time = ros::Time::now();
    // if(current_time.toSec() - global_pose.header.stamp.toSec() > __transform_tolerance)
    // {
    //     ROS_WARN_THROTTLE(1.0,
    //                     "Costmap2DROS transform timeout. Current time: %.4f, global_pose stamp: %.4f, tolerance: %.4f",
    //                     current_time.toSec(), global_pose.header.stamp.toSec(), __transform_tolerance);
    //     return false;
    // }

    return true;   
}

// 获取机器人当前位姿
bool ControllerBase::getRobotPose(tf2_ros::Buffer* tf_buffer, geometry_msgs::PoseStamped& global_pose){
    // 尝试3次
    for(int i=0; i<3; i++){
        if(getTFPose(tf_buffer, global_pose, "map", baseFrameId_, ros::Time::now())){
            return true;
        }
        ros::Duration(0.05).sleep();
    }
    return false;
}


// // 全局位姿转局部
// geometry_msgs::PoseStamped ControllerBase::convertGlobalPoseToLocal(geometry_msgs::PoseStamped global_pose, 
//     geometry_msgs::PoseStamped robot_pose){
//     geometry_msgs::PoseStamped local_pose;

//     // 方式一
//     tf2::Transform global_pose_tf;
//     tf2::fromMsg(global_pose.pose, global_pose_tf);
//     tf2::Transform robot_pose_tf;
//     tf2::fromMsg(robot_pose.pose, robot_pose_tf);

//     // 计算相对变换（全局变换到局部变换）
//     tf2::Transform local_pose_tf = robot_pose_tf.inverse() * global_pose_tf;
//     tf2::toMsg(local_pose_tf, local_pose.pose);
    
//     // 方式二
//     // double yaw = tf2::getYaw(robot_pose.pose.orientation);
//     // double delta_x = global_pose.pose.position.x - robot_pose.pose.position.x;
//     // double delta_y = global_pose.pose.position.y - robot_pose.pose.position.y;
//     // double cos_yaw = std::cos(-yaw);
//     // double sin_yaw = std::sin(-yaw);
//     // local_pose.pose.position.x = cos_yaw * delta_x - sin_yaw * delta_y;
//     // local_pose.pose.position.y = sin_yaw * delta_x + cos_yaw * delta_y;
//     // double yaw_diff = utils::math::diffAngle(
//     //     msg_convert::toYaw(robot_pose.pose.orientation),
//     //     msg_convert::toYaw(global_pose.pose.orientation)
//     // );
//     // local_pose.pose.orientation = msg_convert::toGeometryQuaternion(yaw_diff);

//     return local_pose; 
// }

mc::TrajectoryPoint2D ControllerBase::convertGlobalPoseToLocal(mc::TrajectoryPoint2D global_pose, mc::TrajectoryPoint2D robot_pose){
    mc::TrajectoryPoint2D pose = trajectoryROS_.calculateRelativeCoordinate(global_pose, robot_pose);
    return pose;
}

// 局部位姿转全局
// geometry_msgs::PoseStamped ControllerBase::convertLocalPoseToGlobal(geometry_msgs::PoseStamped local_pose, 
//     geometry_msgs::PoseStamped robot_pose){
//     geometry_msgs::PoseStamped global_pose;

//     // 方式一
//     tf2::Transform local_pose_tf;
//     tf2::fromMsg(local_pose.pose, local_pose_tf);
//     tf2::Transform robot_pose_tf;
//     tf2::fromMsg(robot_pose.pose, robot_pose_tf);

//     // 计算全局位姿（局部位姿在机器人坐标系中的全局变换）
//     tf2::Transform global_pose_tf = robot_pose_tf * local_pose_tf;
//     tf2::toMsg(global_pose_tf, global_pose.pose);

//     // 方式二
//     // double yaw = tf2::getYaw(robot_pose.pose.orientation);
//     // double delta_x = local_pose.pose.position.x;
//     // double delta_y = local_pose.pose.position.y;
//     // double cos_yaw = std::cos(yaw);
//     // double sin_yaw = std::sin(yaw);
//     // global_pose.pose.position.x = robot_pose.pose.position.x + cos_yaw * delta_x - sin_yaw * delta_y;
//     // global_pose.pose.position.y = robot_pose.pose.position.y + sin_yaw * delta_x + cos_yaw * delta_y;
//     // double yaw_add = utils::math::addAngle(
//     //     msg_convert::toYaw(robot_pose.pose.orientation),
//     //     msg_convert::toYaw(local_pose.pose.orientation)
//     // );
//     // global_pose.pose.orientation = msg_convert::toGeometryQuaternion(yaw_add);

//     return global_pose;
// }

mc::TrajectoryPoint2D ControllerBase::convertLocalPoseToGlobal(mc::TrajectoryPoint2D global_pose, mc::TrajectoryPoint2D robot_pose){
    mc::TrajectoryPoint2D pose = trajectoryROS_.calculateRelativeCoordinate(robot_pose, global_pose);
    return pose;
}

/*************************************************
* 辅助函数
*************************************************/
// 旋转至目标角度
// void ControllerBase::rotateToTGoal(geometry_msgs::PoseStamped robot_pose,
//     geometry_msgs::PoseStamped goal_pose, geometry_msgs::Twist& cmd_vel){
//     cmd_vel.linear.x = 0;
//     cmd_vel.linear.y = 0;

//     double angle_diff = utils::math::diffAngle(
//         msg_convert::toYaw(robot_pose.pose.orientation),
//         msg_convert::toYaw(goal_pose.pose.orientation)
//     );
//     // 最大减速 v^2 = v0^2 + 2as
//     double max_speed_to_stop = maxSpeedToStop(baseConfig_.max_vel_yaw, angle_diff);
//     max_speed_to_stop = angle_diff > 0 ? max_speed_to_stop : -max_speed_to_stop;

//     // 最大加速限制 [__min_vel_limit ~ __max_vel_limit]
//     cmd_vel.angular.z = min(baseConfig_.max_vel_yaw, max(baseConfig_.min_vel_yaw, max_speed_to_stop));
// }

void ControllerBase::rotateToTGoal(mc::TrajectoryPoint2D robot_pose,
    mc::TrajectoryPoint2D goal_pose, geometry_msgs::Twist& cmd_vel){
    cmd_vel.linear.x = 0;
    cmd_vel.linear.y = 0;

    double angle_diff = utils::math::diffAngle(
        robot_pose.pose.rotation.angle(), 
        goal_pose.pose.rotation.angle()
    );
    // 最大减速 v^2 = v0^2 + 2as
    double max_speed_to_stop = maxSpeedToStop(baseConfig_.max_vel_yaw, angle_diff);
    max_speed_to_stop = angle_diff > 0 ? max_speed_to_stop : -max_speed_to_stop;

    // 最大加速限制 [__min_vel_limit ~ __max_vel_limit]
    cmd_vel.angular.z = min(baseConfig_.max_vel_yaw, max(baseConfig_.min_vel_yaw, max_speed_to_stop));    
}

// 计算最大减速距离
double ControllerBase::maxSpeedToStop(double vel, double acc){
    // 最大减速 v^2 = v0^2 + 2as
    return sqrt(2 * vel * fabs(acc));
}

// 计算直线刹车距离 v^2 - 0 = 2as
double ControllerBase::maxStopDistance(geometry_msgs::Twist cmd_vel, double acc){
    return std::pow(cmd_vel.linear.x, 2) / (2 * acc);
}

// 检测目标点是否到达  距离判断
// bool ControllerBase::isPositionReached(geometry_msgs::PoseStamped robot_pose, geometry_msgs::PoseStamped goal_pose){
//     double dist = utils::math::distanceBetweenPoses(goal_pose.pose, robot_pose.pose); 
//     LOG(INFO) << "[isPositionReached] dist: " << dist << ", goal_distance_tolerance: " << baseConfig_.goal_distance_tolerance;
//     if(fabs(dist) < baseConfig_.goal_distance_tolerance){
//         return true;
//     }
//     return false;  
// }

bool ControllerBase::isPositionReached(mc::TrajectoryPoint2D robot_pose, mc::TrajectoryPoint2D goal_pose){
    double dist = trajectoryROS_.calculateDistance(goal_pose, robot_pose);
    LOG(INFO) << "[isPositionReached] dist: " << dist << ", goal_distance_tolerance: " << baseConfig_.goal_distance_tolerance;
    if(fabs(dist) < baseConfig_.goal_distance_tolerance){
        return true;
    }
    return false;  
}

// 检测目标点是否到达  角度判断
// bool ControllerBase::isAngleReached(geometry_msgs::PoseStamped robot_pose, geometry_msgs::PoseStamped goal_pose){
    // double theta = utils::math::diffAngle(
    //     msg_convert::toYaw(robot_pose.pose.orientation),
    //     msg_convert::toYaw(goal_pose.pose.orientation)
    // );

    // LOG(INFO) << "[isAngleReached] theta: " << theta << ", goal_angle_tolerance: " << baseConfig_.goal_angle_tolerance;
    // if(fabs(theta) < baseConfig_.goal_angle_tolerance){
    //     return true;
    // }
    // return false;
// }

bool ControllerBase::isAngleReached(mc::TrajectoryPoint2D robot_pose, mc::TrajectoryPoint2D goal_pose){
    double theta = utils::math::diffAngle(
        robot_pose.pose.rotation.angle(), 
        goal_pose.pose.rotation.angle()
    );

    LOG(INFO) << "[isAngleReached] theta: " << theta << ", goal_angle_tolerance: " << baseConfig_.goal_angle_tolerance;
    if(fabs(theta) < baseConfig_.goal_angle_tolerance){
        return true;
    }
    return false;
}

// 检测目标点是否到达
// bool ControllerBase::isGoalReached(geometry_msgs::PoseStamped robot_pose, geometry_msgs::PoseStamped goal_pose){
//     return isPositionReached(robot_pose, goal_pose) && isAngleReached(robot_pose, goal_pose);
// }

bool ControllerBase::isGoalReached(mc::TrajectoryPoint2D robot_pose, mc::TrajectoryPoint2D goal_pose){
    return isPositionReached(robot_pose, goal_pose) && isAngleReached(robot_pose, goal_pose);
}

/*************************************************
* 路径处理
*************************************************/   
// 轨迹裁剪
void ControllerBase::pruningPath(double dis){
    // 从头找到最近点
    size_t closest_index = -1;
    for (size_t i = 0; i < trajectory_points_.size(); ++i) {
        double dist = trajectoryROS_.calculateDistance(trajectory_points_[i], robot_pose_);
        if (dist < dis) {
            closest_index = i;
            break;
        }
    } 

    // 裁剪距离dis米以前的坐标点
    if(closest_index > 0 && closest_index < trajectory_points_.size()){
        trajectory_points_.erase(trajectory_points_.begin(), trajectory_points_.begin() + closest_index - 1);
    }
}   

// 获取前瞻距离
double ControllerBase::getLookAheadDistance(double vt){
    double lookahead_dist = fabs(vt) * baseConfig_.lookahead_time;
    return max(baseConfig_.min_lookahead_distance, 
        min(baseConfig_.max_lookahead_distance, lookahead_dist));
}

// 获取前瞻目标点 简单的距离计算，后期优化
bool ControllerBase::getLookaheadGoal(double lookahead_distance, mc::TrajectoryPoint2D& lookahead_goal, int& index){
    if(trajectory_points_.empty()){
        return false;
    }

    double lookahead_distance_sq = lookahead_distance * lookahead_distance;
    // 寻找后侧最近点
    int closest_index = -1;
    for (size_t i = 0; i < trajectory_points_.size(); ++i) {
        double dis = trajectoryROS_.calculateDistanceSquared(trajectory_points_[i], robot_pose_);
        if(dis < lookahead_distance_sq) {
            closest_index = i;
            break;
        }
    }

    // 沿着索引前进，找到前侧最近点
    if(closest_index >=0){
        for(size_t i = closest_index; i < trajectory_points_.size(); ++i){
            double dis = trajectoryROS_.calculateDistanceSquared(trajectory_points_[i], robot_pose_);
            if(dis > lookahead_distance_sq || i == trajectory_points_.size()-1) {
                lookahead_goal = trajectory_points_[i];
                index = i;
                return true;
            }
        }
    }
    return false;
}


// 获取距离最近的点
bool ControllerBase::getClosedPose(int& closest_index){
    // 设定为最大值
    double dis_sqrt = std::numeric_limits<float>::infinity();
    closest_index = -1;
    for (size_t i = 0; i < trajectory_points_.size(); ++i) {
        double dis = trajectoryROS_.calculateDistanceSquared(trajectory_points_[i], robot_pose_);
        if (dis < dis_sqrt) {
            dis_sqrt = dis;
            closest_index = i;
        }
    }
    return closest_index != -1;
}



// 校验当前点到目标点之间是否存在障碍物
bool ControllerBase::isPathCollision(int& collision_index){
    // TODO: 通过costmap判断路径上是否存在障碍物
    int closest_index = -1;
    // 计算最近的点
    if(getClosedPose(closest_index)){
        // 向前搜索avoidance_distance距离，是否存在障碍物
        double avoidance_distance_sqr = baseConfig_.avoidance_distance * baseConfig_.avoidance_distance;
        // 遍历点集，查看是否存在障碍物
        // auto costmap_ptr = costmapLocalPtr_->getCostmap2D();
        
        for(size_t i = 0; i < trajectory_points_.size(); ++i) {
            double dis = trajectoryROS_.calculateDistanceSquared(trajectory_points_[i], robot_pose_);
            if( dis > avoidance_distance_sqr){
                break;
            }
            else{
                unsigned int mx, my;
                map_base_.worldToMap(trajectory_points_[i].pose.translation.x(), 
                    trajectory_points_[i].pose.translation.y(), mx, my);
                if(!map_base_.valueIsValid(mx, my)){
                    unsigned char cost = map_base_.getCost(mx, my);
                    LOG(INFO) << "i: " << i << ", x: " << trajectory_points_[i].pose.translation.x() 
                        << ", y: " << trajectory_points_[i].pose.translation.y()
                        << ", dis: " << dis
                        << ", mx: " << mx << ", my: " << my  << ", cost: " << int(cost);
                    collision_index = i;
                    return true;
                }
            }
        }
    }

    return false;    
}

// 路径重规划
bool ControllerBase::replanPath(int collision_index){
    // 寻找下个无碰撞点索引和后续无碰撞坐标
    int next_no_collision_index = -1;
    int valid_points_nums = 0;
    for(size_t i = collision_index; i < trajectory_points_.size(); ++i){
        unsigned int mx, my;
        map_base_.worldToMap(trajectory_points_[i].pose.translation.x(), 
            trajectory_points_[i].pose.translation.y(), mx, my);
        if(map_base_.valueIsValid(mx, my)){
            if(next_no_collision_index == -1){
                next_no_collision_index = i;
            }
            else{
                valid_points_nums++;
                if(valid_points_nums > 20){     //最多延展20个点
                    break;  
                }
            }  
        }
        else{

        }
    }
    LOG(INFO) << "next_no_collision_index: " << next_no_collision_index << ", valid_points_nums: " << valid_points_nums;

    if(next_no_collision_index > 0 && valid_points_nums > 5){
        // 重新规划路径
        geometry_msgs::PoseStamped new_goal_pose 
            = trajectoryROS_.toROSPoseStamped(trajectory_points_[next_no_collision_index + valid_points_nums]);
        vector<geometry_msgs::Pose> poses;
        mc::status status = path_planner_ptr_->genPath(trajectoryROS_.toROSPose(robot_pose_), new_goal_pose.pose, poses);
        if(!status.success || poses.empty()){
            LOG(WARNING) << "genPath failed!";
            return false;
        }   
        
        // 拼接路径并平滑处理
        trajectory_points_.erase(trajectory_points_.begin(), trajectory_points_.begin() + next_no_collision_index);
        for(auto point : trajectory_points_){
            poses.emplace_back(trajectoryROS_.toROSPose(point));
        }
        
        // 平滑路并生成路径方向
        vector<geometry_msgs::Pose> smooth_poses;
        path_planner_ptr_->pathSmooth(poses, smooth_poses);
        smooth_poses = path_planner_ptr_->genPathDirection(smooth_poses, poses.back());   // 生成路径方向

        // 更新全局路径
        trajectory_points_.clear();
        trajectory_points_.resize(smooth_poses.size());
        for(int i=0; i<smooth_poses.size(); ++i){
            trajectory_points_[i] = trajectoryROS_.toTrajectoryPoint2D(smooth_poses[i], "map");
        }
        return true;
    }

    return false; 
}

/*************************************************
* 速度控制
*************************************************/
// 急停控制
void ControllerBase::e_stop(){
    sendEmptyTwist();
}

// 减速控制
void ControllerBase::slow_down(double controller_frequency){
    // v = v0 - at
    geometry_msgs::Twist cmd_vel;
    cmd_vel.linear.x = max(0.0, cmd_vel_.linear.x - (baseConfig_.linear_acc * controller_frequency));
    cmd_vel.linear.y = max(0.0, cmd_vel_.linear.y - (baseConfig_.linear_acc * controller_frequency));
    cmd_vel.angular.z = max(0.0, cmd_vel_.angular.z - (baseConfig_.angular_acc * controller_frequency));
    sendVelocity(cmd_vel);
}

// 停止控制
void ControllerBase::sendEmptyTwist(){
    geometry_msgs::Twist empty_twist;
    cmd_vel_ = empty_twist;          //当前速度数据缓存
    sendVelocity(empty_twist);
}

// 发送速度
void ControllerBase::sendVelocity(geometry_msgs::Twist& cmd_vel){
    // 需要添加限幅和平滑处理

    // 速度限制
    cmd_vel.linear.x = max(min(baseConfig_.max_vel_x, cmd_vel.linear.x), baseConfig_.min_vel_x);   
    cmd_vel.linear.y = 0.0;
    cmd_vel.angular.z = max(min(baseConfig_.max_vel_yaw, cmd_vel.angular.z), baseConfig_.min_vel_yaw);

    // cout << "[sendVelocity] cmd_vel: " << msg_convert::toString(cmd_vel) << endl;
    cmd_vel_ = cmd_vel;              //当前速度数据缓存
    speedPub_.publish(cmd_vel);
}